package ecdsa

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"gitee.com/zackeus/go-boot/tools/errorx"
	"github.com/go-jose/go-jose/v4"
)

// GenerateECDSA 生成 ECDSA 密钥对
func GenerateECDSA(alg jose.SignatureAlgorithm) (*ecdsa.PrivateKey, *ecdsa.PublicKey, error) {
	var curve elliptic.Curve
	switch alg {
	case jose.ES256:
		curve = elliptic.P256()
	case jose.ES384:
		curve = elliptic.P384()
	case jose.ES512:
		curve = elliptic.P521()
	}
	privKey, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		return nil, nil, errorx.Wrap(err, "生成 ECDSA 密钥失败")
	}

	return privKey, &privKey.PublicKey, nil
}

// GenerateECDSAPem 生成 ECDSA 密钥对
func GenerateECDSAPem(alg jose.SignatureAlgorithm) ([]byte, []byte, error) {
	privKey, pubKey, err := GenerateECDSA(alg)
	if err != nil {
		return nil, nil, err
	}

	/* 编码私钥 */
	privBytes, err := x509.MarshalECPrivateKey(privKey)
	if err != nil {
		return nil, nil, errorx.Wrap(err, "编码 ECDSA 私钥失败")
	}
	privPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: privBytes,
	})

	/* 编码公钥 */
	pubBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, nil, errorx.Wrap(err, "编码 ECDSA 公钥失败")
	}
	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PUBLIC KEY",
		Bytes: pubBytes,
	})

	return privPEM, pubPEM, nil
}

// LoadECDSAPrivateKey 将 PEM 格式的 ECDSA 私钥加载为 *ecdsa.PrivateKey
func LoadECDSAPrivateKey(privPEM []byte) (*ecdsa.PrivateKey, error) {
	// 解码 PEM 块
	block, _ := pem.Decode(privPEM)
	if block == nil {
		return nil, errors.New("failed to decode PEM block containing private key")
	}

	// 检查 PEM 块类型是否为 EC PRIVATE KEY
	if block.Type != "EC PRIVATE KEY" {
		return nil, errors.New("invalid PEM block type: expected EC PRIVATE KEY, got " + block.Type)
	}

	// 解析 EC 私钥
	privKey, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return nil, errorx.Wrap(err, "failed to parse EC private key")
	}

	return privKey, nil
}

// LoadECDSAPublicKey 将 PEM 格式的 ECDSA 公钥加载为 *ecdsa.PublicKey
func LoadECDSAPublicKey(pubPEM []byte) (*ecdsa.PublicKey, error) {
	// 解码 PEM 块
	block, _ := pem.Decode(pubPEM)
	if block == nil {
		return nil, errors.New("failed to decode PEM block containing public key")
	}

	// 检查 PEM 块类型是否为 EC PUBLIC KEY（与生成时的类型对应）
	if block.Type != "EC PUBLIC KEY" {
		return nil, errors.New("invalid PEM block type: expected EC PUBLIC KEY, got " + block.Type)
	}

	// 解析 PKIX 格式的公钥（x509.MarshalPKIXPublicKey 对应反向操作）
	pubKeyInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, errorx.Wrap(err, "failed to parse EC public key")
	}

	// 将解析出的通用公钥接口断言为 ECDSA 公钥类型
	pubKey, ok := pubKeyInterface.(*ecdsa.PublicKey)
	if !ok {
		return nil, errors.New("parsed public key is not an ECDSA public key")
	}

	return pubKey, nil
}
